#!/usr/bin/env python3
import pandas as pd
import glob
import os

def sum_total_forward_time():
    """
    计算所有timings_all_setting_rlhf_our.csv文件中total_forward_batch_spec_time_ms列的总和
    """
    # 查找所有timings_all_setting_rlhf_our.csv文件
    csv_files = glob.glob("outputs/timings_all_setting_rlhf_*.csv")
    
    def extract_number(filename):
        try:
            return int(filename.split('/')[0].split('_')[1])
        except (ValueError, IndexError):
            return float('inf')  # 将非数字的目录排在最后
    
    csv_files.sort(key=extract_number)
    
    print(f"找到 {len(csv_files)} 个timings_all_setting_rlhf_our.csv文件")
    
    total_sum = 0.0
    file_count = 0
    
    for csv_file in csv_files:
        try:
            # 读取CSV文件
            df = pd.read_csv(csv_file)
            
            # 检查是否存在total_forward_batch_spec_time_ms列
            if 'total_forward_batch_spec_time_ms' in df.columns:
                # 计算该列的总和
                file_sum = df['total_forward_batch_spec_time_ms'].sum()
                total_sum += file_sum
                file_count += 1
                
                # 提取目录编号用于显示
                dir_num = csv_file.split('/')[0].split('_')[1]
                print(f"outputs_{dir_num}: {file_sum:.2f} ms")
            else:
                print(f"警告: {csv_file} 中没有找到 total_forward_batch_spec_time_ms 列")
                
        except Exception as e:
            print(f"处理文件 {csv_file} 时出错: {e}")
    
    print(f"\n总共处理了 {file_count} 个文件")
    print(f"所有文件 total_forward_batch_spec_time_ms 的总和: {total_sum:.2f} ms")
    print(f"平均每个文件: {total_sum/file_count:.2f} ms" if file_count > 0 else "没有有效文件")

if __name__ == "__main__":
    sum_total_forward_time()